import cv2
import numpy as np
import matplotlib.pyplot as plt

# 加载图像
image = cv2.imread('piano512.PNG', cv2.IMREAD_GRAYSCALE)
watermark = cv2.imread('lsc_sign512.jpg', cv2.IMREAD_GRAYSCALE)


# 执行DCT变换
dct_image = cv2.dct(np.float32(image))
dct_watermark = cv2.dct(np.float32(watermark))
# 将水印嵌入到DCT系数中
alpha = 0.01  # 调整水印强度
watermarked_dct = dct_image + alpha * dct_watermark
# 执行逆DCT变换
watermarked_image = cv2.idct(watermarked_dct).real


# 旋转攻击
rotated_watermarked_image = cv2.rotate(watermarked_image, cv2.ROTATE_90_CLOCKWISE)
# 裁剪攻击
cropped_watermarked_image = watermarked_image[50:200, 50:200]
# 调整裁剪攻击后的图像大小
cropped_watermarked_image = cv2.resize(cropped_watermarked_image, (512, 512))
# 添加高斯噪声攻击
noise = np.random.normal(0, 25, watermarked_image.shape).astype(np.float32)
noisy_watermarked_image = watermarked_image + noise


# 原图提取水印
extracted_dct_watermark = (watermarked_dct - dct_image) / alpha
extracted_watermark = cv2.idct(extracted_dct_watermark).real
# 旋转攻击提取水印
rotated_dct_watermarked_image = cv2.dct(np.float32(rotated_watermarked_image))
rotated_extracted_dct_watermark = (rotated_dct_watermarked_image - dct_image) / alpha
rotated_extracted_watermark = cv2.idct(rotated_extracted_dct_watermark).real
# 裁剪攻击提取水印
cropped_dct_watermarked_image = cv2.dct(np.float32(cropped_watermarked_image))
cropped_extracted_dct_watermark = (cropped_dct_watermarked_image - dct_image) / alpha
cropped_extracted_watermark = cv2.idct(cropped_extracted_dct_watermark).real
# 高斯噪声攻击提取水印
noisy_dct_watermarked_image = cv2.dct(np.float32(noisy_watermarked_image))
noisy_extracted_dct_watermark = (noisy_dct_watermarked_image - dct_image) / alpha
noisy_extracted_watermark = cv2.idct(noisy_extracted_dct_watermark).real


# 计算原始图像和带水印图像之间的均方误差 (MSE)
mse_embed = np.mean((dct_image - watermarked_dct) ** 2)
print('嵌入图像和原图 MSE = ', mse_embed)
# 计算SNR
signal_power = np.mean(dct_image ** 2)
noise_power = mse_embed
snr_embed = 10 * np.log10(signal_power / noise_power)
print('嵌入图像 SNR = ', snr_embed)
# 计算旋转攻击后的均方误差和SNR
mse_rotate = np.mean((dct_image - cv2.dct(rotated_watermarked_image)) ** 2)
snr_rotate = 10 * np.log10(signal_power / mse_rotate)
print('旋转攻击后图像 MSE = ', mse_rotate)
print('旋转攻击后图像 SNR = ', snr_rotate)
# 计算裁剪攻击后的均方误差和SNR
mse_crop = np.mean((dct_image - cv2.dct(cropped_watermarked_image)) ** 2)
snr_crop = 10 * np.log10(signal_power / mse_crop)
print('裁剪攻击后图像 MSE = ', mse_crop)
print('裁剪攻击后图像 SNR = ', snr_crop)
# 计算添加高斯噪声攻击后的均方误差和SNR
mse_noise = np.mean((dct_image - cv2.dct(noisy_watermarked_image)) ** 2)
snr_noise = 10 * np.log10(signal_power / mse_noise)
print('添加噪声攻击后图像 MSE = ', mse_noise)
print('添加噪声攻击后图像 SNR = ', snr_noise)


# 显示图像
plt.subplot(4, 3, 1)
plt.imshow(image, cmap='gray')
plt.title('Original Image')
plt.axis('off')

plt.subplot(4, 3, 2)
plt.imshow(watermark, cmap='gray')
plt.title('Watermark')
plt.axis('off')

plt.subplot(4, 3, 3)
plt.imshow(watermarked_image, cmap='gray')
plt.title('Watermarked Image')
plt.axis('off')

plt.subplot(4, 3, 4)
plt.imshow(rotated_watermarked_image, cmap='gray')
plt.title('Rotated Watermarked Image')
plt.axis('off')

plt.subplot(4, 3, 5)
plt.imshow(cropped_watermarked_image, cmap='gray')
plt.title('Cropped Watermarked Image')
plt.axis('off')

plt.subplot(4, 3, 6)
plt.imshow(noisy_watermarked_image, cmap='gray')
plt.title('Noisy Watermarked Image')
plt.axis('off')

plt.subplot(4, 3, 7)
plt.imshow(extracted_watermark, cmap='gray')
plt.title('Extracted Watermark')
plt.axis('off')

plt.subplot(4, 3, 8)
plt.imshow(rotated_extracted_watermark, cmap='gray')
plt.title('rotated_extracted_watermark')
plt.axis('off')

plt.subplot(4, 3, 9)
plt.imshow(cropped_extracted_watermark, cmap='gray')
plt.title('cropped_extracted_watermark')
plt.axis('off')

plt.subplot(4, 3, 10)
plt.imshow(noisy_extracted_watermark, cmap='gray')
plt.title('noisy_extracted_watermark')
plt.axis('off')

plt.show()